import torch


# 简单的卷积神经网络
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model = torch.nn.Sequential(
            #
            # layer 1
            #
            # 输入图片大小为 32x32
            torch.nn.Conv2d(
                in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1
            ),
            # 使用 ReLU 作为激活函数
            torch.nn.ReLU(),
            # 定义一个最大池化层，核大小为 2，步长为 2
            torch.nn.MaxPool2d(kernel_size=2, stride=2),
            #
            # layer 2
            #
            # 经过第一个卷积层和最大池化层后，图片大小减半，变为 16x16
            torch.nn.Conv2d(
                in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1
            ),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2),
            #
            # layer 3
            #
            # 经过第二个卷积层和最大池化层后，图片大小减半，变为 8x8
            torch.nn.Conv2d(
                in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1
            ),
            torch.nn.ReLU(),
            # 将特征图展平为一维向量
            torch.nn.Flatten(),
            #
            # 全连接层
            #
            # 输入特征数为 8 * 8 * 64，输出特征数为 128
            torch.nn.Linear(in_features=8 * 8 * 64, out_features=128),
            torch.nn.ReLU(),
            #
            # 输出层
            #
            # 输入特征数为 128，输出特征数为 10（对应0-9共十个种类）
            torch.nn.Linear(in_features=128, out_features=10),
            # 使用 Softmax 函数进行概率分布，使得输出的每一项都在 0-1 之间，并且和为 1
            torch.nn.Softmax(dim=1),
        )

    def forward(self, input):
        # 把张量传入Sequential，得到前向传播的结果
        output = self.model(input)
        return output
